from functools import cached_property, reduce
import re
from .seq import MTSequence, Sequence
from .gene import Transcript, RegionType
from .snv import SNV, TranscriptAnnotation


class SUB(SNV):
    @cached_property
    def s_start(self) -> int:
        return self.pos

    @cached_property
    def s_ref(self) -> Sequence | MTSequence:
        return self.ref

    @cached_property
    def s_alt(self) -> Sequence | MTSequence:
        """standard alt"""
        return self.alt

    @cached_property
    def size1(self) -> int:
        return self.ref.size

    @cached_property
    def size2(self) -> int:
        return self.alt.size
    
    def set_aa_change(self, trans_anno: TranscriptAnnotation, transcript: Transcript, pos: int) -> TranscriptAnnotation:
        cdna = transcript.coding_dna
        new_cdna = cdna.substitute(pos, self.ref, self.alt)
        if transcript.strand == "-":
            cdna = cdna.reverse_complement()
            new_cdna = new_cdna.reverse_complement()
        cstart = cdna.difference_simple(new_cdna)
        trans_anno.na_change = f"c.{cstart}_{cstart+self.size1-1}delins{self.alt}"
        protein = cdna.protein()
        new_protein = new_cdna.protein()
        pstart, pend1, pend2 = protein.difference(new_protein)
        if pstart > protein.size:
            # 同义变异
            pstart = ((cstart - 1) // 3) + 1
            aa1, aa2 = protein[pstart - 1], new_protein[pstart - 1]
            trans_anno.aa_change = f"p.{aa1}{pstart}{aa2}"
            trans_anno.event = "sub_synonymous"
        else:
            # 非同义变异
            aa1 = protein[pstart - 1 : pend1]
            aa2 = new_protein[pstart - 1 : pend2]
            if (cdna.size - new_cdna.size) % 3 == 0:
                trans_anno.event = "sub_inframe"
                if aa2.size == 0:
                    trans_anno.aa_change = f"p.{aa1}{pstart}del" if aa1.size == 1 else f"p.{aa1[0]}{pstart}_{aa1[-1]}{pend1}del"
                else:
                    trans_anno.aa_change = f"p.{aa1}{pstart}delins{aa2}" if aa1.size == 1 else f"p.{aa1[0]}{pstart}_{aa1[-1]}{pend1}delins{aa2}"
            else:
                if pstart < protein.size:
                    trans_anno.event = "sub_frameshift"
                    if aa2[0] == "*":
                        trans_anno.aa_change = f"p.{aa1[0]}{pstart}*"
                    else:
                        fsi = new_protein[pstart - 1 :].find("*")
                        fs = "?" if fsi == -1 else ("" if fsi == 0 else f"{fsi+1}")
                        trans_anno.aa_change = f"p.{aa1[0]}{pstart}{aa2[0]}fs*{fs}"
            if protein == "M" and new_protein and new_protein[0] != "M":
                trans_anno.event += "_startloss"
            if trans_anno.region.find("splic") != -1:
                trans_anno.event += "_splicing"

    def annotate_mrna(self, transcript: Transcript) -> TranscriptAnnotation:
        region1, cds_size1, utr_size1 = transcript.find_region(self.s_start)
        region2, cds_size2, utr_size2 = transcript.find_region(self.s_end)
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region1, region2)

        if region1 != region2:
            if region2.start - region1.end > 1:
                trans_anno.region = "deletion"
            else:
                if {region1.type, region2.type} == {RegionType.CDS, RegionType.INTRON}:
                    trans_anno.region = "exonic_splicing"
        if (self.s_start < transcript.cds_start or self.s_start >= transcript.cds_end) and (self.s_end < transcript.cds_start or self.s_end >= transcript.cds_end):
            """
            snv包含了整个编码区，即整个编码区被删除 SNV的起始终止位置都不在CDS区或CDS区之间的Intron中
            情况2: snv包含了整个编码区，即整个编码区被删除
              ...+++,,,+++...
                |---------|
             |---------------|
            情况2: snv发生在整个编码区左边
              |-|...+++,,,+++...
                |-|
            情况3: snv发生在整个编码区左边
             ...+++,,,+++...|-|
                          |-|
            """
            distance1 = transcript.calc_utr_distance(region1, utr_size1, self.s_start)
            distance2 = transcript.calc_utr_distance(region2, utr_size2, self.s_end)
            trans_anno.na_change = f"c.{distance1}_{distance2}delins{self.alt}" if transcript.strand == "+" else f"c.{distance2}_{distance1}delins{self.alt.reverse_complement()}"
            if self.s_start < transcript.cds_start and self.s_end >= transcript.cds_end:
                # 情况1
                trans_anno.region = "transcript"
                trans_anno.region2 = "transcript"
                trans_anno.event = "deletion"
        elif self.s_start < transcript.cds_start <= self.s_end <= transcript.cds_end:
            """
             ...+++,,,+++...
              |--|
            |----|
            """
            distance1 = transcript.calc_utr_distance(region1, utr_size1, self.s_start)
            distance2 = transcript.calc_cds_distance(region2, cds_size2, self.s_end)
            trans_anno.na_change = f"c.{distance1}_{distance2}delins{self.alt}" if transcript.strand == "+" else f"c.{distance2}_{distance1}delins{self.alt.reverse_complement()}"
        elif transcript.cds_start <= self.s_start <= transcript.cds_end < self.s_end:
            """
             ...,,,+++...
                    |--|
                    |----|
            """
            distance1 = transcript.calc_cds_distance(region1, cds_size1, self.s_start)
            distance2 = transcript.calc_utr_distance(region2, utr_size2, self.s_end)
            trans_anno.na_change = f"c.{distance1}_{distance2}delins{self.alt}" if transcript.strand == "+" else f"c.{distance2}_{distance1}delins{self.alt.reverse_complement()}"
        else:
            """"变异跨region"""
            distance1 = transcript.calc_cds_distance(region1, cds_size1, self.s_start)
            distance2 = transcript.calc_cds_distance(region2, cds_size2, self.s_end)
            if region1 != region2:
                trans_anno.na_change = f"c.{distance1}_{distance2}delins{self.alt}" if transcript.strand == "+" else f"c.{distance2}_{distance1}delins{self.alt.reverse_complement()}"
            else:
                if region1.type == RegionType.INTRON:
                    trans_anno.na_change = f"c.{distance1}_{distance2}delins{self.alt}" if transcript.strand == "+" else f"c.{distance2}_{distance1}delins{self.alt.reverse_complement()}"
                    values = list(map(int, reduce(lambda x, y: x + y, map(lambda z: re.findall(r"[\+\-](\d+)$", z), [distance1, distance2]))))
                    if values and min(values) <= 2:
                        trans_anno.event = "splicing"
                        trans_anno.region = "splicing"
                else:
                    pos = cds_size1 + (self.s_start - region1.start + 1)
                    self.set_aa_change(trans_anno, transcript, pos)
        return trans_anno
    def annotate_ncrna(self, transcript: Transcript) -> TranscriptAnnotation:
        trans_anno = TranscriptAnnotation.from_transcript(transcript)
        trans_anno.region = "ncRNA"
        pos = self.s_start - transcript.tx_start + 1
        dna = transcript.dna
        new_dna = dna.substitute(pos, self.ref, self.alt)
        if transcript.strand == "-":
            dna = dna.reverse_complement()
            new_dna = new_dna.reverse_complement()
        start = new_dna.difference_simple(dna)
        alt = self.alt if transcript.strand == "+" else self.alt.reverse_complement()
        trans_anno.na_change = f"n.{start}_{start+self.size1-1}delins{alt}"
        return trans_anno
